package com.github.filter;

import com.github.util.AESEncrypt;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

@Component
public class EncryptFilter extends OncePerRequestFilter {
    private static final List<String> FILTER_URI = new ArrayList<>();

    static {
        FILTER_URI.add("/api/server_login");
        FILTER_URI.add("/api/server_logout");
        FILTER_URI.add("/api/get_setting");
    }

    @Override
    protected boolean shouldNotFilter(HttpServletRequest request) throws ServletException {
        String requestURI = request.getRequestURI();
        return !requestURI.startsWith("/api/");
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain)
            throws ServletException, IOException {
        String path = request.getRequestURI();
        if (!FILTER_URI.contains(path)) {
            filterChain.doFilter(request, response);
            return;
        }
        // 使用自定义的HttpServletResponseWrapper来捕获响应体
        CharResponseWrapper responseWrapper = new CharResponseWrapper(response);
        responseWrapper.setContentType("application/json; charset=UTF-8");

        // 继续处理请求
        filterChain.doFilter(request, responseWrapper);

        // 获取响应体
        String responseBody = new String(responseWrapper.getContentAsByteArray(), StandardCharsets.UTF_8);

        // 加密
        String data = AESEncrypt.encrypt(responseBody);
        String payload = "{\"payload\":\"" + data + "\"}";

        // 将加密后的响应体写回原始响应
        responseWrapper.resetBuffer();
        responseWrapper.getWriter().write(payload);
        responseWrapper.getWriter().flush();
        responseWrapper.copyBodyToResponse();
    }

    public static class CharResponseWrapper extends HttpServletResponseWrapper {
        private final ByteArrayOutputStream content = new ByteArrayOutputStream();
        private PrintWriter writer;
        private ServletOutputStream outputStream;

        public CharResponseWrapper(HttpServletResponse response) {
            super(response);
        }

        @Override
        public ServletOutputStream getOutputStream() throws IOException {
            if (outputStream == null) {
                outputStream = new ServletOutputStream() {
                    @Override
                    public boolean isReady() {
                        return true;
                    }

                    @Override
                    public void setWriteListener(WriteListener writeListener) {
                        throw new UnsupportedOperationException("Not supported");
                    }

                    @Override
                    public void write(int b) throws IOException {
                        content.write(b);
                    }
                };
            }
            return outputStream;
        }

        public byte[] getContentAsByteArray() {
            if (writer != null) {
                writer.flush();
            }
            if (outputStream != null) {
                try {
                    outputStream.flush();
                } catch (IOException e) {
                    throw new RuntimeException("Failed to flush output stream", e);
                }
            }
            return content.toByteArray();
        }

        public void copyBodyToResponse() throws IOException {
            byte[] bytes = getContentAsByteArray();
            if (bytes.length > 0) {
                getResponse().setContentLength(bytes.length);
                getResponse().getOutputStream().write(bytes);
            }
        }

        @Override
        public PrintWriter getWriter() throws IOException {
            if (writer == null) {
                writer = new PrintWriter(new OutputStreamWriter(content, getCharacterEncoding()));
            }
            return writer;
        }

        public void resetBuffer() {
            content.reset();
        }
    }
}
